import math
from collections import Counter


def distance(p1, p2):
    return math.sqrt(sum((p1[i] - p2[i]) ** 2 for i in range(len(p1))))


def get_k_neighbors(train_X, train_y, test_x, k):
    dists = []
    for i in range(len(train_X)):
        d = distance(test_x, train_X[i])
        dists.append((d, train_y[i]))

    dists.sort(key=lambda x: x[0])

    return dists[:k]


# KNN 分类
def KNN_classify(train_X, train_y, test_x, k=3):
    neighbors = get_k_neighbors(train_X, train_y, test_x, k)
    labels = [label for _, label in neighbors]

    most_common = Counter(labels).most_common(1)[0][0]
    return most_common


# KNN 回归
def KNN_regress(train_X, train_y, test_x, k=3):
    neighbors = get_k_neighbors(train_X, train_y, test_x, k)
    values = [label for _, label in neighbors]

    return sum(values) / len(values)


def KNN_predict(train_X, train_y, test_x, k=3, mode="classify"):
    if mode == "classify":
        return KNN_classify(train_X, train_y, test_x, k)
    elif mode == "regress":
        return KNN_regress(train_X, train_y, test_x, k)
    else:
        raise ValueError("mode must be 'classify' or 'regress'")
